#include "ReadsMapping.h"
#include "MappingResult.h"

const int MISQUERYCHECKPOINT = 100000;

CReadsMapping::CReadsMapping(void)
{
    this->initialization();
}

CReadsMapping::CReadsMapping(const MappingOpts P)
{
    this->initialization();
    this->opt= P;

    if (P.bGetAllAlignments) {
        // WARNING if index indexTable.bExcludeAmbiguous will have some error
        alignmentsQ[0].setQueue_All_Best_OneFlag('A');
        alignmentsQ[1].setQueue_All_Best_OneFlag('A');
    }

    if (P.outputFormat != '\0') {
        this->cOutputFormat = P.outputFormat[0];
    }

    // check if the director exist, if not, create one
    if (P.outputDir[0] != '\0') {
        if (dirExist(P.outputDir) || createdir(P.outputDir) == 0) {
#ifdef  WIN32
            sprintf(this->opt.outputDir, "%s\\", P.outputDir);
#else
            sprintf(this->opt.outputDir, "%s/", P.outputDir);
#endif
        } else {
            LOG_INFO("Info %d: Can't create dir %s.\n", WARNING_LOG, P.outputDir);
            this->opt.outputDir[0] = '\0';
        }
    }
}

CReadsMapping::~CReadsMapping(void)
{
    delete AlignResult;
    delete MissReads;
}

void CReadsMapping::initialization(void)
{
    this->AlignResult = NULL;
    this->MissReads = NULL;
    this->cOutputFormat = 's';
}

int CReadsMapping::dealMissedRead(CMappingResult& m)
{
    if (m.QScores[0] == '\0') {
        sprintf(MissReads->caBufp, ">%s\n%s\n", m.QNAME, m.caRead);
        MissReads->UpdateSize();
    } else {
        sprintf(MissReads->caBufp, "@%s\n%s\n+\n%s\n", m.QNAME, m.caRead, m.QScores);
        MissReads->UpdateSize();
    }
    return(this->iMissReadCounter++);
}

// This private function returns a string as the prefix of the output file name for mapping
string getOutputFileNamePrefix(const char* dir, const CGenome_Index_TableQ& table, MappingOpts& opt)
{
    // The prefix contains 4 parts: (0) dir (1) ref file (2) ambiguous flag (3) mismatch threshold
    char outputPrefixStr[MAX_PATH];
    char excludeAmbiguous = opt.bExcludeAmbiguousReads ? 'E' : 'B';
    if (opt.bGetAllAlignments) excludeAmbiguous = 'A';
    sprintf(outputPrefixStr, "%s%s_%c_%u_%d_%d",\
            dir, table.caRefName, excludeAmbiguous, table.chosenSeedId, opt.subDiffThreshold, opt.anchorLength);
    return(string(outputPrefixStr));
}

bool isGalaxyOutputPath(const char* fileName)
{
    return(hasTheExtName(fileName, ".dat"));
}

bool isSupportedExtName(const char* fileN)
{
    if (hasTheExtName(fileN, ".fasta") ||
            hasTheExtName(fileN, ".fastq") ||
            hasTheExtName(fileN, ".fq") ||
            hasTheExtName(fileN, ".csfq") ||
            hasTheExtName(fileN, ".csfastq") ||
            hasTheExtName(fileN, ".csfasta") ||
            hasTheExtName(fileN, ".dat") ||
            hasTheExtName(fileN, ".txt")) {
        return(true);
    }
    return(false);
}

string rmSupportedExtName(const char* fileN)
{
    char newFileN[FILENAME_MAX];
    strcpy(newFileN, fileN);
    if (isSupportedExtName(fileN)) {
        for (int i = (int)strlen(newFileN); i > 0; i--) {
            if (newFileN[i] == '.') {
                newFileN[i] = '\0';
                return(string(newFileN));
            }
        }
    }
    return(string(fileN));
}

// generate the unmapped FileN, if the original one is not null
void getUnmappedFileN(char* unmappedFileN, const char* readFileN, const char* readFormat,\
                      const char* refN, bool bQual = false)
{
    string refNStr = rmSupportedExtName(refN);
    string readFileNStr = rmSupportedExtName(readFileN);
    if (strcmp(unmappedFileN, "") == 0) { // make sure the readFormat is valid
        if (strcmp(readFormat, "csfasta") == 0) {
            if (bQual) {
                sprintf(unmappedFileN, "%s_miss_%s.%s", readFileNStr.c_str(), refNStr.c_str(), "fastq");
            } else {
                sprintf(unmappedFileN, "%s_miss_%s", readFileNStr.c_str(), refNStr.c_str()); // Not sure csfasta or not
            }
        } else {
            if (readFormat[0] == '\0') {
                sprintf(unmappedFileN, "%s_miss_%s", readFileNStr.c_str(), refNStr.c_str());
            } else {
                sprintf(unmappedFileN, "%s_miss_%s.%s", readFileNStr.c_str(), refNStr.c_str(), readFormat);
            }
        }
    }
}

int CReadsMapping::setUpIO4Aligment(const char* caReadsSetName, const CGenome_Index_TableQ& table)
{
    if (this->opt.bGetAllAlignments) {
        // Find all alignments with in threshold, instead of queuing
        // only all best(with the fewest mismatches) alignments as default.
        alignmentsQ[0].setQueue_All_Best_OneFlag('A');
        alignmentsQ[1].setQueue_All_Best_OneFlag('A');
    }
    char outputPath[MAX_PATH];
    if (this->opt.outputFileN[0] == '\0') { // No output file has been set
        string outFileNPrefix = getOutputFileNamePrefix(this->opt.outputDir, table, opt);
        string fileNameOfReadSet = getBasename(caReadsSetName);
        sprintf(outputPath, "%s_%s", outFileNPrefix.c_str(), fileNameOfReadSet.c_str());
    } else {
        sprintf(outputPath, "%s%s", this->opt.outputDir, this->opt.outputFileN);
    }

    // (1)Initialize the I/O for output the alignments
    if (this->opt.bPrintAlignments) {
        if (isGalaxyOutputPath(this->opt.outputFileN)) {
            // Don't change the ext name for Galaxy.
        } else if (this->cOutputFormat == 'm') {
            chExtName(outputPath, ".mapping");
        } else {
            chExtName(outputPath, ".sam");
        }
        if (this->AlignResult == NULL) {
            ofstream* AlignResultFile = new ofstream(outputPath);
            this->AlignResult =
                new FileOutputBuffer(ALIGNMENT_RESULT_FILE_BUFFER_SIZE, AlignResultFile);
            if (this->AlignResult == NULL) {
                ERR;//check new FileOutputBuffer
                return(1);
            } else if (this->cOutputFormat == 's') {
                if (this->opt.bPrintSamHeader) {
                    string RG = getSamRG(caReadsSetName, table.bMapReadInColors);
                    vector<CGene> refs = table.pgenomeNT->getRefNamesLengths();
                    string commandLineStr = this->opt.fullCommand;
                    printSamHeader(AlignResult, refs, RG.c_str(), commandLineStr.c_str());
                }
            }
        }
    }
    // (2)Initialize the I/O for output the up-mapped reads.
    if (this->opt.bPrintUnMappedReads) {
        // Make sure opt.readsFileFormat has been set
        const char* refN = table.pgenomeNT->refName;
        getUnmappedFileN(opt.unmappedFileN, caReadsSetName, opt.readsFileFormat, refN);
        if (!isGalaxyOutputPath(opt.unmappedFileN)) {
            string extName = string(".").append(string(opt.readsFileFormat));
            chExtName(opt.unmappedFileN, extName.c_str());
        }
        if (this->MissReads == NULL) {
            ofstream* MissReadsFile = new ofstream(opt.unmappedFileN);
            this->MissReads =
                new FileOutputBuffer(ALIGNMENT_RESULT_FILE_BUFFER_SIZE, MissReadsFile);
            if (this->MissReads == NULL) {
                ERR; // check new FileOutputBuffer
                return(1);
            }
        }
    }
    return(0);
}

int CReadsMapping::tearDownIO4Aligment(void)
{
    // close file, which has been opened before
    if (this->opt.bPrintAlignments) {
        // this->AlignResult->removeEndBlankLine();
        delete this->AlignResult;
        this->AlignResult = NULL;
    }
    if (this->opt.bPrintUnMappedReads) {
        // this->MissReads->removeEndBlankLine();
        delete this->MissReads;
        this->MissReads = NULL;
    }
    return(0);
}

void getSingleMappingSeq4Solexa(const CGenome_Index_TableQ& table, CMappingResult& m, bool samFormat)
{
    if (samFormat) {
        m.caRef[0] = '\0';
    } else { // sam format doesn't need to know reference sequence
        const unsigned int readLength = (unsigned int)strlen(m.caRead);
        // (1) Get ref in bits
        CReadInBits ref = table.pgenomeNTInBits->getSubstringInBits(m.uiGlobalMappedPos, readLength);
        if (m.strand == '-') {
            ref = reverseCompliment(readLength, ref);
        }
        // (2) Get ref seq
        ref.decode(m.caRef);
    }
}

void getReadQscores4Solexa(CAlignmentsQ& aQue, CMappingResult& m, bool samFormat)
{
    aQue.read.decode(m.caRead);
    if (aQue.qualityScores == NULL) {
        m.mismatchScore = (int)m.uiDiff;
        m.QScores[0] = '\0';
    } else { // If quality score are available, get quality score
        m.mismatchScore = alignmentScore(m.caRead, m.caRef, m.uiReadLength, aQue.qualityScores);
        trQScores(m.uiReadLength, SolexaScoreEncodingShift, aQue.qualityScores, m.QScores);
        if (samFormat) {
            m.getReverseReadandQual();
        }
    }
}

void getSingleMappingSeqAndQ4SOLiD\
(const CGenome_Index_TableQ& table, CAlignmentsQ& aQue, CMappingResult& m, bool samFormat)
{
    const unsigned int readlength = table.uiRead_Length;
    // (1) Get ref in bits
    CReadInBits ref = table.pgenomeNTInBits->getSubstringInBits(m.uiGlobalMappedPos, readlength);
    if (m.strand == '-') {
        ref = reverseCompliment(readlength, ref);
    }
    // (2) Get ref (and read) in Seq
    decodeColors(m.caRead, aQue.read);
    CReadInBits colorsInRef = bases2Colors(ref);
    // (3) Get nucleotide and quality sequence
    if (samFormat) { // correctRead
        if (aQue.qualityScores == NULL) {
            char dummyScore = (char)(20);
            fillDummyQScores(readlength, dummyScore + Phred_SCALE_QUAL_SHIFT, m.QScores);
        } else {
            trQScores(readlength, Phred_SCALE_QUAL_SHIFT, aQue.qualityScores, m.QScores);
        }
        sprintf(m.TAG, "%s\tCS:Z:%s\tCQ:Z:%s", m.TAG, m.caRead, m.QScores);
        //sprintf(m.TAG, "%s\tX1:Z:%s\tCS:Z:%s\tCQ:Z:%s", m.TAG, m.caRef, m.caRead, m.QScores);
        correctAndDecodeRead(aQue.read, colorsInRef, samFormat, m.caRead, m.QScores);
        if (m.strand == '-')
            m.getReverseReadandQual();
    } else {
        decodeColors(m.caRef, colorsInRef); // Show the reference for validation purpose
        m.SNPtype = returnSNPtype(aQue.read, colorsInRef);
        if (aQue.qualityScores == NULL) {
            m.mismatchScore = (int)m.uiDiff;
            m.QScores[0] = '\0';
            m.rawScores[0] = '\0';
        } else {
            m.mismatchScore = alignmentScore(m.caRead, m.caRef,\
                                             readlength, aQue.qualityScores);
            printCommaSepScoresStr(readlength, aQue.qualityScores, m.QScores);
        }
    }
}

// For one read to one location,
void getSingleMappingInfo(const CGenome_Index_TableQ& table, CAlignmentsQ& aQue,\
                          unsigned int mappingId, CMappingResult& m, bool samFormat)
{
    getSingleMappingIndex(*(table.pgenomeNT), aQue, mappingId, m);
    if (samFormat) { // must set after the strand flag is set
        m.setSingleEndSamFields();
    } // must set before the color seq and QUAL is set
    if (table.bMapReadInColors) {
        getSingleMappingSeqAndQ4SOLiD(table, aQue, m, samFormat);
    } else {
        getSingleMappingSeq4Solexa(table, m, samFormat);
    }
    if (table.pbaRepeatRepresentativeFlag->b(aQue.aiHitIndex[mappingId])) {
        m.MultipleMappedNo++;
    }
}

// TODO change this to get indel mapping info
void getLongMappingInfo(const CGenome_Index_TableQ& table, CAlignmentsQ& aQue, bool samFormat,\
                        unsigned int mappingId, CMappingResult& m)
{
    // (1) Get index
    getSingleMappingIndex(*(table.pgenomeNT), aQue, mappingId, m);
    if (samFormat) {
        m.setSingleEndSamFields();
    } else {
        // (2) Get Sequence
        CReadInBits ref1stHalf = table.pgenomeNTInBits->getSubstringInBits\
                                 (m.uiGlobalMappedPos, table.uiRead_Length);
        CReadInBits ref2ndHalf = table.pgenomeNTInBits->getSubstringInBits\
                                 (m.uiGlobalMappedPos + table.uiRead_Length, table.uiRead_Length);
        ref1stHalf.decode(m.caRef);
        ref2ndHalf.decode(&(m.caRef[table.uiRead_Length]));
        // (3) Get mismatch score
        if (m.QScores[0] == '\0') {
            m.mismatchScore = (int)m.uiDiff;
        } else { // If quality score are available, get quality score
            if (m.strand == '-') {
                reverseComplementKmer(m.caRef);
            }
            m.mismatchScore = alignmentScore(m.caRead, m.caRef, m.uiReadLength, m.rawScores);
        }
    }
}

bool wrongIndex(const CReadInBitsSet& readsSet, const CGenome_Index_TableQ& table)
{
    if (readsSet.uiRead_Length != table.uiRead_Length) {
        LOG_INFO("Info %d: The index is not for read length %d, is for %d.\n",\
                 readsSet.uiRead_Length, table.uiRead_Length, ERROR_LOG);
        return(true);
    }
    bool bColorRead = (readsSet.cFileType == 'Q' || readsSet.cFileType == 'S');
    if ( bColorRead != table.bMapReadInColors) {
        if (bColorRead) {
            LOG_INFO("Info %d: The index is not for color read.\n", ERROR_LOG);
        } else {
            LOG_INFO("Info %d: The index is for color read.\n", ERROR_LOG);
        }
        return(true);
    }
    return(false);
}
